from selenium.webdriver.edge.webdriver import WebDriver
from PIL import Image
from io import BytesIO
import cv2
import sys
import time


class Verification:
    SCALING_UP = 1.25
    SCALING_DOWN = 0.75
    JS_CODE_DICT = {
        'no_gap': '''document.getElementsByClassName('geetest_canvas_fullbg')[0].style.display="block";''',
        'with_gap': '''document.getElementsByClassName('geetest_canvas_fullbg')[0].style.display="none";'''
    }
    FAIL_FLAGS = ["正确拼合", "请重试"]
    RECONNECTION_CHANCE = 2
    DISTANCE_MODIFY = [5, -5]

    def __init__(self, driver: WebDriver):
        self.driver = driver
        if not len(self.DISTANCE_MODIFY) == self.RECONNECTION_CHANCE:
            print("验证重试次数和距离修正次数不相等！")
            sys.exit(1)

    def get_screenshot(self):
        return Image.open(BytesIO(self.driver.get_screenshot_as_png()))

    def __get_yzm_img(self, location: dict, size: dict):
        time.sleep(5)
        screenshot = self.get_screenshot()
        # 问: 传入crop的tuple元素报类型不对
        # element_area = [location['x'], location['y'],
        #                 location['x'] + size['width'],
        #                 location['y'] + size['height']]
        # screenshot.crop(tuple(map(lambda item: item * self.SCALING_UP, element_area)))
        yzm_img = screenshot.crop((location['x'] * self.SCALING_UP, location['y'] * self.SCALING_UP,
                                   (location['x'] + size['width']) * self.SCALING_UP,
                                   (location['y'] + size['height']) * self.SCALING_UP))
        return yzm_img

    def get_yzm_img_no_gap(self, location: dict, size: dict):
        # 通过js将图片上方的缺损和验证图块消去，得到完整图片
        self.driver.execute_script(self.JS_CODE_DICT['no_gap'])
        time.sleep(3)
        # 获取当前验证码图片
        yzm_img = self.__get_yzm_img(location, size)
        yzm_img.save(".\\gap_images\\no_gap.png")
        # 通过js对于验证码区域进行复原
        self.driver.execute_script(self.JS_CODE_DICT['with_gap'])
        time.sleep(3)

    def get_yzm_img_with_gap(self, location: dict, size: dict):
        # 获取带有缺口的验证码图片
        yzm_img = self.__get_yzm_img(location, size)
        yzm_img.save(".\\gap_images\\with_gap.png")

    def find_slide_distance(self):
        try:
            img_no_gap = cv2.imread(".\\gap_images\\no_gap.png")
            img_with_gap = cv2.imread(".\\gap_images\\with_gap.png")
        except FileNotFoundError:
            print("no_gap.png/with_gap.png is not found!")
            sys.exit(1)
        # 框选待移动目标拼图区域，可被优化为固定阈值
        bbox = cv2.selectROI("select_gap_result", img_with_gap)
        cv2.destroyAllWindows()
        gap_area_list = self.__find_gap_area(bbox, img_no_gap, img_with_gap)
        # 返回需要滑动的距离
        return int(gap_area_list[0][0]) * self.SCALING_DOWN

    @staticmethod
    def __find_gap_area(bbox, img_no_gap, img_with_gap):
        area_list = list()
        # 高、宽、通道数
        for x in range(bbox[2], img_with_gap.shape[1]):
            for y in range(0, img_with_gap.shape[0]):
                pixel1 = img_with_gap[y, x]
                pixel2 = img_no_gap[y, x]
                for i in range(0, 3):
                    if abs(int(pixel1[i]) - int(pixel2[i])) >= 60:
                        area_list.append((x, y))
                        break
        # 找到图片中的缺口处，填充蓝色并保存
        for pixel_point in area_list:
            cv2.circle(img_with_gap, pixel_point, 1, (255, 0, 0), 1)
        cv2.imwrite(".\\gap_images\\gap_result.png", img_with_gap)
        # 返回可能是缺口的像素点列表
        return area_list

    def is_verify_fail(self, result_content: str):
        is_fail = False
        for flag in self.FAIL_FLAGS:
            if flag in result_content:
                is_fail = True
                break
        return is_fail

    def distance_modify(self, chance, distance):
        # 进行距离修正
        return distance + self.DISTANCE_MODIFY[chance]

    @staticmethod
    def check_yzm_format(input_yzm: str):
        # 格式正确时返回True
        return len(input_yzm) == 6 and input_yzm.isdigit()
